# -*- coding: utf-8 -*-
"""
成交量对比筛选器
筛选条件：
1. 当前成交量仅为最近30天最高成交量的三分之一及以下
2. 当前5日平均成交量低于10日平均成交量（缩量趋势）

使用方式示例：
    python volume_ratio_screener.py
    python volume_ratio_screener.py --limit 200
    python volume_ratio_screener.py --stock 000001
    python volume_ratio_screener.py --days 40  # 自定义比较天数
"""
import argparse
import logging
import os
import sys
from datetime import datetime
from typing import Optional

import pandas as pd

# 添加父目录到路径
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(PROJECT_ROOT)

from config import DATABASE_CONFIG  # noqa: E402
from database import DatabaseManager, StockInfoDAO  # noqa: E402


LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'volume_ratio_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class VolumeRatioScreener:
    """成交量对比筛选器"""

    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
        self.stock_dao = StockInfoDAO(self.db)

    def get_stock_codes(self, limit: Optional[int] = None):
        sql = (
            "SELECT stock_code, stock_name "
            "FROM new_stock_info "
            "WHERE status = 'active' "
            "ORDER BY stock_code"
        )
        if limit:
            sql += f" LIMIT {limit}"

        try:
            rows = self.db.execute_query(sql)
            return [(row["stock_code"], row["stock_name"]) for row in rows]
        except Exception as exc:
            logger.error("获取股票列表失败: %s", exc)
            return []

    def get_stock_daily_data(self, stock_code: str, days: int):
        """获取指定天数的日线数据（按日期升序）"""
        sql = """
        SELECT trade_date, close_price, volume
        FROM new_stock_daily
        WHERE stock_code = %s
        ORDER BY trade_date DESC
        LIMIT %s
        """
        try:
            rows = self.db.execute_query(sql, (stock_code, days))
            if not rows:
                return pd.DataFrame()
            df = pd.DataFrame(rows)
            df["trade_date"] = pd.to_datetime(df["trade_date"])
            df["close_price"] = pd.to_numeric(df["close_price"], errors="coerce")
            df["volume"] = pd.to_numeric(df["volume"], errors="coerce")
            return df.sort_values("trade_date").reset_index(drop=True)
        except Exception as exc:
            logger.error("获取股票 %s 数据失败: %s", stock_code, exc)
            return pd.DataFrame()

    def check_conditions(self, df: pd.DataFrame, compare_days: int = 30):
        """检查是否满足成交量对比与均量关系"""
        min_required_days = max(compare_days, 10)
        if df.empty or len(df) < min_required_days:
            return None
        if df["volume"].isna().any():
            return None

        recent_df = df.tail(compare_days)
        latest_row = recent_df.iloc[-1]
        latest_volume = latest_row["volume"]

        max_volume = recent_df["volume"].max()
        if max_volume <= 0:
            return None

        volume_ratio = latest_volume / max_volume
        if volume_ratio > (1 / 3):
            return None

        ma5 = df["volume"].tail(5).mean()
        ma10 = df["volume"].tail(10).mean()
        if pd.isna(ma5) or pd.isna(ma10):
            return None
        if ma5 >= ma10:
            return None

        return {
            "trade_date": latest_row["trade_date"],
            "close_price": latest_row["close_price"],
            "latest_volume": round(latest_volume, 2),
            "max_volume": round(max_volume, 2),
            "volume_ratio": round(volume_ratio, 4),
            "ma5_volume": round(ma5, 2),
            "ma10_volume": round(ma10, 2),
            "compare_days": compare_days,
        }

    def screen_stocks(self, stock_codes=None, limit: Optional[int] = None, compare_days: int = 30):
        if stock_codes is None:
            stock_codes = self.get_stock_codes(limit=limit)
        if not stock_codes:
            logger.warning("没有需要筛选的股票")
            return pd.DataFrame()

        logger.info("开始筛选 %d 只股票（对比最近%d天）...", len(stock_codes), compare_days)
        results = []
        total = len(stock_codes)
        fetch_days = max(compare_days, 10)

        for idx, (stock_code, stock_name) in enumerate(stock_codes, start=1):
            if idx % 10 == 0 or idx == total:
                logger.info("进度: %d/%d (%.1f%%)", idx, total, idx / total * 100)

            df = self.get_stock_daily_data(stock_code, days=fetch_days)
            if df.empty:
                continue

            result = self.check_conditions(df, compare_days=compare_days)
            if not result:
                continue

            result["stock_code"] = stock_code
            result["stock_name"] = stock_name
            results.append(result)

            logger.info(
                "✓ %s %s - 收: %.2f 成交量比例: %.2f%% MA5: %.0f MA10: %.0f",
                stock_code,
                stock_name,
                result["close_price"],
                result["volume_ratio"] * 100,
                result["ma5_volume"],
                result["ma10_volume"],
            )

        if not results:
            logger.info("未找到符合条件的股票")
            return pd.DataFrame()

        df_results = pd.DataFrame(results)
        df_results = df_results.sort_values("volume_ratio")
        return df_results[
            [
                "stock_code",
                "stock_name",
                "trade_date",
                "close_price",
                "latest_volume",
                "max_volume",
                "volume_ratio",
                "ma5_volume",
                "ma10_volume",
                "compare_days",
            ]
        ]

    def save_results(self, df: pd.DataFrame, output_dir: str = "."):
        if df.empty:
            return None, None
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_file = os.path.join(output_dir, f"volume_ratio_stocks_{timestamp}.csv")
        txt_file = os.path.join(output_dir, f"stock_codes_volume_ratio_{timestamp}.txt")
        df.to_csv(csv_file, index=False, encoding="utf-8-sig")
        with open(txt_file, "w", encoding="utf-8") as file:
            for _, row in df.iterrows():
                file.write(f"{row['stock_code']}\n")
        logger.info("结果已保存至: %s", csv_file)
        logger.info("股票代码已保存至: %s", txt_file)
        return csv_file, txt_file

    def print_summary(self, df: pd.DataFrame):
        if df.empty:
            return
        logger.info("符合条件股票数量: %d", len(df))
        logger.info("平均成交量比例: %.2f%%", df["volume_ratio"].mean() * 100)
        logger.info("最小成交量比例: %.2f%%", df["volume_ratio"].min() * 100)


def parse_args():
    parser = argparse.ArgumentParser(description="成交量对比筛选器")
    parser.add_argument("--limit", type=int, help="限制筛选股票数量")
    parser.add_argument("--stock", type=str, help="指定单只股票代码")
    parser.add_argument("--days", type=int, default=30, help="比较天数（默认为30天）")
    parser.add_argument("--output-dir", type=str, default=".", help="输出目录")
    return parser.parse_args()


def main():
    args = parse_args()

    screener = VolumeRatioScreener()
    if args.stock:
        stock_info = screener.stock_dao.get_stock_info_by_code(args.stock)
        if not stock_info:
            logger.error("未找到股票: %s", args.stock)
            return
        stock_codes = [(stock_info["stock_code"], stock_info["stock_name"])]
        logger.info("筛选单只股票: %s %s", stock_info["stock_code"], stock_info["stock_name"])
    else:
        stock_codes = None

    start_time = datetime.now()
    df = screener.screen_stocks(stock_codes=stock_codes, limit=args.limit, compare_days=args.days)
    end_time = datetime.now()

    screener.print_summary(df)
    if not df.empty:
        screener.save_results(df, output_dir=args.output_dir)
        logger.info("筛选耗时: %.2f 秒", (end_time - start_time).total_seconds())

    logger.info("筛选完成")


if __name__ == "__main__":
    main()

